package com.gobang.websocket;

import com.gobang.common.exception.BusinessException;
import com.gobang.common.service.TokenAuthService;
import com.gobang.common.service.TokenAuthService.AuthUser;
import com.gobang.websocket.service.ReconnectTokenService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.util.UriComponentsBuilder;

import java.net.URI;
import java.util.List;
import java.util.Map;

/**
 * WebSocket 握手拦截器
 *
 * 在握手阶段完成 JWT + Redis 双重校验，并将 userId、username、role 等信息写入 attributes。
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class JwtHandshakeInterceptor implements HandshakeInterceptor {

    private static final String AUTHORIZATION_HEADER = "Authorization";
    private static final String TOKEN_PREFIX = "Bearer ";

    private final TokenAuthService tokenAuthService;
    private final ReconnectTokenService reconnectTokenService;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) {
        try {
            String token = extractToken(request);
            if (!StringUtils.hasText(token)) {
                log.warn("WebSocket 握手失败：Token 缺失");
                return false;
            }

            AuthUser authUser = tokenAuthService.validateToken(token);
            Long userId = authUser.getUserId();
            attributes.put("userId", userId);
            attributes.put("username", authUser.getUsername());
            attributes.put("role", authUser.getRole());

            // 处理断线重连令牌：若存在则进行校验
            String reconnectToken = extractReconnectToken(request);
            if (StringUtils.hasText(reconnectToken)) {
                String roomId = reconnectTokenService.validate(reconnectToken, userId);
                attributes.put("reconnectToken", reconnectToken);
                attributes.put("roomId", roomId);
            }

            return true;
        } catch (BusinessException e) {
            log.warn("WebSocket 握手鉴权失败: code={}, message={}", e.getCode(), e.getMessage());
            return false;
        } catch (Exception e) {
            log.error("WebSocket 握手鉴权异常", e);
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        // 无需处理
    }

    private String extractToken(ServerHttpRequest request) {
        // 1. 优先从 Query 参数中获取 token
        URI uri = request.getURI();
        String token = UriComponentsBuilder.fromUri(uri)
                .build()
                .getQueryParams()
                .getFirst("token");

        if (StringUtils.hasText(token)) {
            return token;
        }

        // 2. 兼容 Authorization: Bearer xxx 方式
        HttpHeaders headers = request.getHeaders();
        String authHeader = headers.getFirst(AUTHORIZATION_HEADER);
        if (StringUtils.hasText(authHeader) && authHeader.startsWith(TOKEN_PREFIX)) {
            return authHeader.substring(TOKEN_PREFIX.length());
        }

        // 3. 兼容 Servlet 环境下从 HttpServletRequest 读取（可选）
        if (request instanceof ServletServerHttpRequest servletRequest) {
            String paramToken = servletRequest.getServletRequest().getParameter("token");
            if (StringUtils.hasText(paramToken)) {
                return paramToken;
            }
        }

        return null;
    }

    private String extractReconnectToken(ServerHttpRequest request) {
        URI uri = request.getURI();
        List<String> values = UriComponentsBuilder.fromUri(uri)
                .build()
                .getQueryParams()
                .get("reconnectToken");
        if (values != null && !values.isEmpty()) {
            return values.get(0);
        }
        if (request instanceof ServletServerHttpRequest servletRequest) {
            return servletRequest.getServletRequest().getParameter("reconnectToken");
        }
        return null;
    }
}
